import tkinter as tk
from tkinter import filedialog, messagebox
from PIL import Image, ImageTk
import cv2
import numpy as np
import logging

logging.basicConfig(format='%(asctime)s:%(levelname)s:%(message)s', level=logging.INFO)

class App:
    def __init__(self, root):
        self.root = root
        self.img = None
        self.mask = None
        self.file_path = None
        self.need_refresh = False
        self.current_method = None
        self.prev_pts = []
        self.pen_strokes = []
        self.canvas_width = 350
        self.canvas_height = int(self.canvas_width * 0.75)
        
        self.create_widgets()
        self.bind_brush_action()

    def create_widgets(self):
        # 界面布局优化
        tk.Button(self.root, text="载入图片", command=self.load_image, font=('Arial', 10)).pack(pady=10)
        
        self.image_frame = tk.Frame(self.root)
        self.canvas1 = tk.Canvas(self.image_frame, width=self.canvas_width, height=self.canvas_height, bg="white")
        self.canvas2 = tk.Canvas(self.image_frame, width=self.canvas_width, height=self.canvas_height, bg="white")
        self.canvas1.pack(side=tk.LEFT, padx=5)
        self.canvas2.pack(side=tk.LEFT, padx=5)
        
        # 添加标题标签
        tk.Label(self.image_frame, text="原图", font=('Arial', 10)).place(x=self.canvas_width//2-15, y=-20)
        tk.Label(self.image_frame, text="修复后", font=('Arial', 10)).place(x=self.canvas_width*3//2-15, y=-20)
        
        self.image_frame.pack(pady=20)
        
        button_frame = tk.Frame(self.root)
        tk.Button(button_frame, text="FMM", command=self.apply_fmm, width=8).pack(side=tk.LEFT, padx=5)
        tk.Button(button_frame, text="NS", command=self.apply_ns, width=8).pack(side=tk.LEFT, padx=5)
        tk.Button(button_frame, text="Undo", command=self.undo, width=8).pack(side=tk.LEFT, padx=5)
        tk.Button(button_frame, text="Refresh", command=self.refresh, width=8).pack(side=tk.LEFT, padx=5)
        button_frame.pack(pady=10)
        
        # 关键修复：恢复滑动条的command处理，并确保窗口大小为奇数
        self.slider = tk.Scale(self.root, from_=3, to=51, orient=tk.HORIZONTAL, 
                               length=200, resolution=2, command=self.slider_update, label="窗口大小")
        self.slider.set(7)
        self.slider.pack(pady=10)
        
        # 添加使用说明
        instructions = """
        使用说明:
        1. 点击"载入图片"选择图像
        2. 在左侧原图上用鼠标绘制需要修复的区域（白色线条）
        3. 选择FMM/NS算法后，拖动滑动条调整窗口大小（奇数效果更佳）
        """
        tk.Label(self.root, text=instructions, justify=tk.LEFT, font=('Arial', 9)).pack(pady=5)

    def bind_brush_action(self):
        self.canvas1.bind("<B1-Motion>", self.draw_line)
        self.canvas1.bind("<ButtonRelease-1>", self.reset_prev_point)

    def load_img_mask(self):
        try:
            img_buff = np.fromfile(self.file_path, dtype=np.uint8)
            self.img = cv2.imdecode(img_buff, cv2.IMREAD_COLOR)
            
            if self.img is None:
                raise ValueError("无效的图像文件")
            
            h, w = self.img.shape[:2]
            ratio = min(self.canvas_height / h, self.canvas_width / w)
            self.img = cv2.resize(self.img, (int(w*ratio), int(h*ratio)))
            self.mask = np.zeros(self.img.shape[:2], dtype=np.uint8)
            
        except Exception as e:
            messagebox.showerror("加载错误", f"图像加载失败：{str(e)}")
            self.file_path = None
            self.img = None
            self.mask = None

    def load_image(self):
        self.file_path = filedialog.askopenfilename(filetypes=[
            ("图像文件", "*.png;*.jpg;*.jpeg;*.bmp;*.gif")
        ])
        
        if not self.file_path:
            return
        
        self.pen_strokes = []
        self.canvas2.delete("all")
        self.load_img_mask()
        
        if self.img is not None:
            pil_img = Image.fromarray(cv2.cvtColor(self.img, cv2.COLOR_BGR2RGB))
            self.photo = ImageTk.PhotoImage(pil_img)
            self.canvas1.create_image(0, 0, anchor=tk.NW, image=self.photo)
            logging.info(f"成功加载图像: {self.file_path}")

    def reset_prev_point(self, event):
        if self.prev_pts:
            self.pen_strokes.append(self.prev_pts.copy())
            logging.info(f"保存笔画，总数：{len(self.pen_strokes)}")
        self.prev_pts = []

    def draw_line(self, event):
        if self.img is None:
            return
        
        x, y = event.x, event.y
        if 0 <= x < self.canvas_width and 0 <= y < self.canvas_height:
            pt = (x, y)
            if self.prev_pts:
                pen_width = 5
                cv2.line(self.img, self.prev_pts[-1], pt, (255, 255, 255), pen_width)
                cv2.line(self.mask, self.prev_pts[-1], pt, 255, pen_width)
                self.update_canvas()
            self.prev_pts.append(pt)

    def update_canvas(self):
        if self.img is not None:
            pil_img = Image.fromarray(cv2.cvtColor(self.img, cv2.COLOR_BGR2RGB))
            self.photo = ImageTk.PhotoImage(pil_img)
            self.canvas1.create_image(0, 0, anchor=tk.NW, image=self.photo)

    def apply_fmm(self):
        self.current_method = 'fmm'
        self.run_inpaint(cv2.INPAINT_TELEA)

    def apply_ns(self):
        self.current_method = 'ns'
        self.run_inpaint(cv2.INPAINT_NS)

    def run_inpaint(self, method):
        if self.img is None or self.mask is None:
            messagebox.showwarning("提示", "请先加载图像并标记修复区域")
            return
        
        if np.sum(self.mask) == 0:
            messagebox.showwarning("提示", "请绘制需要修复的区域")
            return
        
        window_size = self.get_odd_window_size()
        try:
            logging.info(f"执行修复: 算法={method}, 窗口大小={window_size}")
            res = cv2.inpaint(self.img, self.mask, window_size, method)
            pil_img = Image.fromarray(cv2.cvtColor(res, cv2.COLOR_BGR2RGB))
            after_photo = ImageTk.PhotoImage(pil_img)
            self.canvas2.create_image(0, 0, anchor=tk.NW, image=after_photo)
            self.canvas2.image = after_photo
            logging.info("图像修复成功")
        except Exception as e:
            messagebox.showerror("处理错误", f"修复失败：{str(e)}")

    def get_odd_window_size(self):
        """获取奇数窗口大小，确保OpenCV正确处理"""
        window_size = int(self.slider.get())
        return window_size if window_size % 2 == 1 else window_size + 1

    def refresh(self):
        self.need_refresh = True
        self.current_method = None
        self.load_image()

    def undo(self):
        if not self.pen_strokes:
            messagebox.showinfo("提示", "无可撤销操作")
            return
        
        self.pen_strokes.pop()
        self.reload_original_image()
        
        for stroke in self.pen_strokes:
            for i in range(len(stroke)-1):
                cv2.line(self.img, stroke[i], stroke[i+1], (255, 255, 255), 5)
                cv2.line(self.mask, stroke[i], stroke[i+1], 255, 5)
        
        self.update_canvas()

    def reload_original_image(self):
        if self.file_path and self.img is not None:
            h, w = self.img.shape[:2]
            self.load_img_mask()
            if self.img is not None:
                self.img = cv2.resize(self.img, (w, h))  # 恢复原始缩放尺寸
                self.mask = np.zeros(self.img.shape[:2], dtype=np.uint8)

    def slider_update(self, val):
        """滑动条值变化时触发，根据当前算法更新修复"""
        if self.current_method == 'fmm':
            self.apply_fmm()
        elif self.current_method == 'ns':
            self.apply_ns()

if __name__ == "__main__":
    root = tk.Tk()
    root.title("图像修复工具")
    root.geometry("800x600")
    app = App(root)
    root.mainloop()